#pragma once

#include "tree.hpp"
#include <map>

template<class T, class W>
struct huffman_node : public binary_node<T>
{
private:
	typedef T data_type;
	typedef W weight;

	typedef binary_node<data_type> base_class;

	typedef huffman_node<data_type, weight> self;
	typedef self* self_pointer;
public:
	explicit huffman_node(const data_type& data = data_type(), const weight& w = weight())
		: base_class(data)
		, _weight(w)
	{}

	virtual void swap(self& right)
	{
		base_class::swap(right);
		std::swap(_weight, right._weight);
	}

	weight _weight = weight();
};

template<class T, class W>
class huffman_tree : public binary_tree<T>
{
	typedef T data_type;
	typedef W weight;

	typedef struct huffman_node<data_type, weight> huffman_node;
	typedef huffman_node* node_pointer;

	typedef binary_tree<data_type> base_class;
	typedef huffman_tree<data_type, weight> self;

	struct queue_cmp
	{
		bool operator()(node_pointer left, node_pointer right)
		{
			return (left->_weight > right->_weight);
		}
	};
public:
	huffman_tree(std::initializer_list<data_type> _Ilist)
		: base_class()
	{
		new(this)self(_Ilist.begin(), _Ilist.end());
	}
	template<class InputIterator>
	huffman_tree(InputIterator first, InputIterator end)
		: base_class()
	{
		std::priority_queue<node_pointer, std::vector<node_pointer>, queue_cmp> heap;
		while (first != end)
		{
			node_pointer node = new huffman_node(first->first, first->second);
			heap.push(node);
			++first;
		}
		while (heap.size() != 1)
		{
			node_pointer left = heap.top();
			heap.pop();
			node_pointer right = heap.top();
			heap.pop();
			node_pointer parent = new huffman_node(data_type(), left->_weight + right->_weight);
			parent->_pLeft = left;
			parent->_pRight = right;
			heap.push(parent);
		}
		base_class::_root = heap.top();
		heap.pop();
		std::vector<bool> prifix;
		encode_map(prifix, dynamic_cast<node_pointer>(base_class::_root));
	}
	
	const std::vector<bool>& compress(const data_type& right) const
	{
		typename std::map<data_type, std::vector<bool>>::const_iterator cit = _map.find(right);
		if (cit != _map.end())
			return cit->second;
		else
			throw std::invalid_argument("compress");
	}

	const data_type uncompress(const std::vector<bool>& right) const
	{
		data_type ret;
		node_pointer cur = dynamic_cast<node_pointer>(base_class::_root);
		for (std::vector<bool>::const_iterator cit = right.begin(); true; ++cit)
		{
			if (cur == nullptr)
				throw std::invalid_argument("uncompress");
			else if (cur->_pLeft == nullptr && cur->_pRight == nullptr)
			{
				if (cur == dynamic_cast<node_pointer>(base_class::_root))
					throw std::invalid_argument("uncompress");
				ret = cur->_data;
				break;
			}
			if (*cit == false)
				cur = dynamic_cast<node_pointer>(cur->_pLeft);
			else if (*cit == true)
				cur = dynamic_cast<node_pointer>(cur->_pRight);
		}
		return ret;
	}

private:
	void encode_map(std::vector<bool>& prifix, node_pointer root)
	{
		if (root == nullptr)
			return;
		if (root->_pLeft == nullptr && root->_pRight == nullptr)
			_map.emplace(root->_data, prifix);
		//L 0	R 1
		if (root->_pLeft != nullptr)
		{
			prifix.push_back(false);
			encode_map(prifix, dynamic_cast<node_pointer>(root->_pLeft));
			prifix.pop_back();
		}
		if (root->_pRight != nullptr)
		{
			prifix.push_back(true);
			encode_map(prifix, dynamic_cast<node_pointer>(root->_pRight));
			prifix.pop_back();
		}
	}

	std::map<data_type, std::vector<bool>> _map;
};

